import math
import os
import json
import math
import pandas as pd
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QColor
from PyQt5.QtWidgets import QMessageBox, QLabel, QLineEdit, QDialog, QGridLayout, QVBoxLayout, QTableWidget, \
    QHeaderView, QTableWidgetItem, QTableView

from utils.math import calculate_kadian_point, calculate_zhonghe_point, calculate_mozuli, calculate_zhouxiangli, \
    calculate_wanquyingli, calculate_ququ_status, fuzhengqi_recommend
from utils.schema import DrillingParameters
from view.custom_comobox import HistoryLineEdit
from view.database_dialog import show_database_viewer
from view.custom_dialog import WellManagerDialog
from utils.huitu3d import ThreeDPlotWidget

from database.db_manager import DatabaseManager


class JieKaDaLao:
    def __init__(self, main):
        self.main = main
        self.ui = main.ui
        self.data_import_view = main.data_import_view
        self.current_well_name = None
        self.data_modified = False
        self.well_data = {}
        self.data_file = "well_data.json"

        # 加载历史数据
        self.load_data()

        # 获取所有字段的描述
        schema = DrillingParameters.model_json_schema()

        self.parameters = [field_info.get('description', '无描述') for field_info in schema['properties'].values()]
        self.input_fields = {}
        # 初始化UI
        self.init_ui()

    def init_ui(self):
        # 钻柱设计-选择钻杆按钮
        layout = QGridLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        self.ui.widget_para_input.setLayout(layout)
        self.create_parameter_form()
        self.ui.pushButton_sel_zg.clicked.connect(self.on_show_database)

        self.ui.pushButton_well.clicked.connect(self.manage_wells)
        self.ui.pushButton_new.clicked.connect(self.new_well)
        self.ui.pushButton_save.clicked.connect(self.save_well)

        self.ui.pushButton_ok.clicked.connect(self.calculater)
        if not self.ui.jkdl_Draw_widget.layout():
            self.ui.jkdl_Draw_widget.setLayout(QVBoxLayout())
        
        # 初始化tableWidget_zzsj表格
        self.init_tableWidget_zzsj()

    def init_tableWidget_zzsj(self):
        """初始化tableWidget_zzsj表格，设置列名并从wells表加载列名"""
        # 设置基本列：自、至、钻杆名称
        basic_columns = ["钻杆名称", "自", "至"]
        # 从wells表获取所有列名
        well_columns = self.get_wells_table_columns()
        
        # 合并列名
        all_columns = basic_columns + well_columns
        
        # 设置表格列数
        self.ui.tableWidget_zzsj.setColumnCount(len(all_columns))
        
        # 设置列标题
        self.ui.tableWidget_zzsj.setHorizontalHeaderLabels(all_columns)
        
        # 设置表格属性
        self.ui.tableWidget_zzsj.setAlternatingRowColors(True)  # 交替行颜色
        self.ui.tableWidget_zzsj.setShowGrid(True)  # 显示网格
        self.ui.tableWidget_zzsj.setSelectionBehavior(QTableWidget.SelectRows)  # 选择整行
        self.ui.tableWidget_zzsj.setSelectionMode(QTableWidget.SingleSelection)  # 单选模式
        
        # 设置列宽
        header = self.ui.tableWidget_zzsj.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.Stretch)  # "钻杆名称"列拉伸
        header.setSectionResizeMode(1, QHeaderView.ResizeToContents)  # "自"列
        header.setSectionResizeMode(2, QHeaderView.ResizeToContents)  # "至"列
        
        # wells表列设置为可调整
        for i in range(3, len(all_columns)):
            header.setSectionResizeMode(i, QHeaderView.Interactive)
        
        # 设置初始行数（可以为0，等待数据填充）
        self.ui.tableWidget_zzsj.setRowCount(0)
        
        # 连接表格信号，当单元格内容改变时调用on_data_modified
        self.ui.tableWidget_zzsj.cellChanged.connect(self.on_table_cell_changed)

    def on_table_cell_changed(self, row, column):
        """表格单元格内容改变时的回调函数"""
        self.on_data_modified()

    def get_wells_table_columns(self):
        """从wells表获取所有列名，使用DatabaseManager"""
        try:
            # 创建数据库管理器实例
            db_manager = DatabaseManager()
            
            # 连接数据库
            if not db_manager.connect():
                print("无法连接到数据库")
                return self._get_default_columns()
            
            try:
                # 使用PRAGMA table_info获取表结构
                query = db_manager.execute_raw_query("PRAGMA table_info(wells)")
                columns = []
                
                while query.next():
                    column_name = query.value(1)  # 第2列是列名
                    columns.append(column_name)
                
                # 确保查询对象在使用后被销毁
                query.finish()
                query = None
                
                if not columns:
                    print("wells表不存在或没有列")
                    return self._get_default_columns()
                
                return columns
                
            finally:
                # 确保断开数据库连接
                db_manager.disconnect()
                
        except Exception as e:
            print(f"获取wells表列名时出错: {str(e)}")
            return self._get_default_columns()
    
    def _get_default_columns(self):
        """获取默认列名"""
        return ["序列号", "公称直径", "外径/mm", "壁厚/mm", "钢级", "名义重量kg/m",
               "长度/m", "水眼直径/mm", "内压强度/MPa", "挤毁强度/MPa",
               "拉伸载荷（屈服）/MPa", "拉伸载荷（最大）/MPa", "抗扭载荷（管体）/MPa",
               "接头类型", "接头外径/mm", "上扣扭矩/N·m", "接头屈服扭矩/N·m",
               "接头抗拉载荷/kN", "适用井深范围/m", "安全系数"]

    def on_show_database(self):
        """点击按钮时显示数据库查看对话框"""
        show_database_viewer(self.ui, self)  # 传入self.ui作为父窗口，self作为JiekadalaoView实例

    def on_show_3d(self, depths, widget_draw, parent):
        """点击按钮时显示卡点计算 3d 图"""
        three_d_plot_widget = ThreeDPlotWidget(self.data_import_view.original_data, widget_draw)
        three_d_plot_widget.show_plot_dialog(parent)

        three_d_plot_widget.set_marked_points(depths)

    def calculater(self):
        """计算方法，获取表单数据并转换为 DrillingParameters 模型后进行计算"""
        try:
            # 检查油井名称
            form_data = self.get_form_data()
            well_name = form_data.get("井号", "").strip()
            if not well_name:
                QMessageBox.warning(None, "警告", "请输入井号!")
                return
            # 获取表单数据并转换为 DrillingParameters 模型
            drilling_params = self.get_drilling_parameters(form_data)
            kadian_depth = calculate_kadian_point(drilling_params)
            zhonghedian_depth = calculate_zhonghe_point(drilling_params)
            depths = [(kadian_depth, "卡点"), (zhonghedian_depth, "中和点")]
            self.on_show_3d(depths, self.ui.jkdl_Draw_widget, self.ui.tab_draw_3D)
            
            # 计算摩阻力
            # 获取井轨迹数据（从data_import_view获取）
            if self.data_import_view.original_data is not None:
                df_origin_data = self.data_import_view.original_data.copy()
                
                # 获取钻杆设计表格数据
                df_zgsj = pd.DataFrame(self.get_table_data())
                
                # 调用摩阻力计算函数
                df_mozuli_result = calculate_mozuli(drilling_params, df_origin_data, df_zgsj)
                
                # 调用轴向力计算函数
                df_zhouxiangli_result = calculate_zhouxiangli(drilling_params, df_mozuli_result, df_zgsj)
                
                # 调用弯曲应力计算函数
                df_wanquyingli_result = calculate_wanquyingli(drilling_params, df_zhouxiangli_result, df_zgsj)
                
                # 调用屈曲状态计算函数
                df_ququ_result = calculate_ququ_status(drilling_params, df_wanquyingli_result, df_zgsj)

                self.display_cal_result(df_ququ_result)
                
                # 调用扶正器推荐
                df_fzq_result = fuzhengqi_recommend(drilling_params, df_ququ_result, df_zgsj)

                self.display_fzq_result(df_fzq_result)

            else:
                QMessageBox.warning(None, "警告", "请先导入井轨迹数据!")
        except ValueError as e:
            QMessageBox.warning(None, "数据验证错误", f"计算失败：{str(e)}")
        except Exception as e:
            QMessageBox.critical(None, "计算错误", f"计算过程中发生错误：{str(e)}")
            
    def display_cal_result(self, df_result):
        """在tableWidget中显示摩阻力计算结果"""
        try:
            # 设置表格的行数和列数
            self.ui.cal_res_tableWidget.setRowCount(len(df_result))
            self.ui.cal_res_tableWidget.setColumnCount(len(df_result.columns))
            
            # 设置表头
            self.ui.cal_res_tableWidget.setHorizontalHeaderLabels(df_result.columns.tolist())
            self.ui.cal_res_tableWidget.setSelectionBehavior(QTableView.SelectRows)
            self.ui.cal_res_tableWidget.setSelectionMode(QTableView.ExtendedSelection)
            self.ui.cal_res_tableWidget.setSortingEnabled(False)  # 禁用排序功能
            self.ui.cal_res_tableWidget.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)  # 这里设置可调整列宽
            # self.ui.cal_res_tableWidget.setAlternatingRowColors(True)  # 启用交替行颜色
            self.ui.cal_res_tableWidget.setShowGrid(True)  # 显示网格
            self.ui.cal_res_tableWidget.setGridStyle(Qt.SolidLine)
            self.ui.cal_res_tableWidget.setEditTriggers(QTableView.NoEditTriggers)  # 禁止直接编辑

            # 填充数据
            for row in range(len(df_result)):
                for col in range(len(df_result.columns)):
                    item = QTableWidgetItem(str(df_result.iloc[row, col]))
                    self.ui.cal_res_tableWidget.setItem(row, col, item)
            
            # 调整列宽
            self.ui.cal_res_tableWidget.resizeColumnsToContents()

            # 如果屈曲状态列的值等于"螺旋屈曲"，则将改行颜色设置为红色
            for row in range(len(df_result)):
                if df_result.iloc[row]["屈曲状态"] == "螺旋屈曲":
                    # 将该整行背景色设置成淡红色
                    # 遍历该行的所有列
                    col_count = self.ui.cal_res_tableWidget.columnCount()
                    for col in range(col_count):
                        self.ui.cal_res_tableWidget.item(row, col).setBackground(QColor(255, 192, 203))

        except Exception as e:
            QMessageBox.critical(None, "显示错误", f"显示计算结果时发生错误：{str(e)}")

    def display_fzq_result(self, df_result):
        """在fzq_tableWidget中显示扶正器推荐结果"""
        try:
            # 设置表格的行数和列数
            self.ui.fzq_tableWidget.setRowCount(len(df_result))
            self.ui.fzq_tableWidget.setColumnCount(len(df_result.columns))
            
            # 设置表头
            self.ui.fzq_tableWidget.setHorizontalHeaderLabels(df_result.columns.tolist())
            self.ui.fzq_tableWidget.setSelectionBehavior(QTableView.SelectRows)
            self.ui.fzq_tableWidget.setSelectionMode(QTableView.ExtendedSelection)
            self.ui.fzq_tableWidget.setSortingEnabled(False)  # 禁用排序功能
            self.ui.fzq_tableWidget.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
            self.ui.fzq_tableWidget.setAlternatingRowColors(True)  # 启用交替行颜色
            self.ui.fzq_tableWidget.setShowGrid(True)  # 显示网格
            self.ui.fzq_tableWidget.setGridStyle(Qt.SolidLine)
            self.ui.fzq_tableWidget.setEditTriggers(QTableView.NoEditTriggers)  # 禁止直接编辑

            # 填充数据
            for row in range(len(df_result)):
                for col in range(len(df_result.columns)):
                    item = QTableWidgetItem(str(df_result.iloc[row, col]))
                    self.ui.fzq_tableWidget.setItem(row, col, item)
            
            # 调整列宽
            # self.ui.fzq_tableWidget.resizeColumnsToContents()
            
        except Exception as e:
            QMessageBox.critical(None, "显示错误", f"显示扶正器推荐结果时发生错误：{str(e)}")

    def create_parameter_form(self):
        gridLayout_input = self.ui.widget_para_input.layout()
        # 定义油井参数字段
        for i, param in enumerate(self.parameters):
            label = QLabel(param + ":")
            input_field = HistoryLineEdit(str(i))
            gridLayout_input.addWidget(label, i, 1)
            gridLayout_input.addWidget(input_field, i, 2)
            self.input_fields[param] = input_field
            input_field.textChanged.connect(self.on_data_modified)

        # 设置第二列（输入框所在列）的最小宽度
        gridLayout_input.setColumnMinimumWidth(2, 150)

    def on_data_modified(self):
        """数据修改标志"""
        self.data_modified = True
        self.update_status()

    def update_status(self):
        """更新状态栏"""
        if self.data_modified:
            self.ui.label_status.setText("有未保存的更改")
            self.ui.label_status.setStyleSheet("color: red;")
        else:
            self.ui.label_status.setText("")

    def load_data(self):
        """从JSON文件加载数据 左边表单数据"""
        if os.path.exists(self.data_file):
            try:
                with open(self.data_file, 'r', encoding='utf-8') as f:
                    self.well_data = json.load(f)
            except Exception as e:
                QMessageBox.warning(None, "警告", f"加载数据失败: {str(e)}")
                self.well_data = {}

    def save_data(self):
        """保存数据到JSON文件"""
        try:
            with open(self.data_file, 'w', encoding='utf-8') as f:
                json.dump(self.well_data, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            QMessageBox.warning(None, "错误", f"保存数据失败: {str(e)}")
            return False

    def get_form_data(self):
        """获取表单数据"""
        data = {}
        for param, field in self.input_fields.items():
            data[param] = field.text()
        return data

    def get_table_data(self):
        """获取tableWidget_zzsj表格数据"""
        table_data = []
        row_count = self.ui.tableWidget_zzsj.rowCount()
        column_count = self.ui.tableWidget_zzsj.columnCount()
        
        for row in range(row_count):
            row_data = {}
            for col in range(column_count):
                header_text = self.ui.tableWidget_zzsj.horizontalHeaderItem(col).text()
                item = self.ui.tableWidget_zzsj.item(row, col)
                if item:
                    row_data[header_text] = item.text()
                else:
                    row_data[header_text] = ""
            table_data.append(row_data)
        
        return table_data

    def set_table_data(self, table_data):
        """设置tableWidget_zzsj表格数据"""
        # 暂时断开信号连接，避免在加载数据时触发on_data_modified
        self.ui.tableWidget_zzsj.cellChanged.disconnect()
        
        try:
            if not table_data:
                self.ui.tableWidget_zzsj.setRowCount(0)
                return
            
            # 设置行数
            self.ui.tableWidget_zzsj.setRowCount(len(table_data))
            column_count = self.ui.tableWidget_zzsj.columnCount()
            
            # 填充数据
            for row_idx, row_data in enumerate(table_data):
                for col_idx, (header_text, value) in enumerate(row_data.items()):
                    if col_idx < column_count:
                        item = QTableWidgetItem(str(value))
                        self.ui.tableWidget_zzsj.setItem(row_idx, col_idx, item)
        finally:
            # 重新连接信号
            self.ui.tableWidget_zzsj.cellChanged.connect(self.on_table_cell_changed)

    def get_drilling_parameters(self, form_data):
        """将表单数据转换为 DrillingParameters 模型实例"""

        # 创建字段映射，将中文描述映射到模型字段名
        field_mapping = {
            "井号": "well_name",
            "井深（m）": "well_depth",
            "修井液密度（kg/m³）": "workover_fluid_density",
            "钢材密度（kg/m³）": "steel_density",
            "弹性模量（N/m²）": "elastic_modulus",
            "摩擦系数": "friction_coefficient",
            "大钩载荷（kN）": "safety_factor",
            "井口压力（MPa）": "tensile_safety_factor",
            "上提拉力（kN）": "pulling_force",
            "钻具伸长度（mm）": "collar_inner_outer_diameter",
            "被卡物体截面积（mm²）": "cross_sectional_area",
            "井筒内径（mm）": "slip_length",
            "修井液动切力（Pa）": "collar_length",
            "修井液塑性粘度（MPa·s）": "collar",
            "泵排量（L/min）": "length"
        }

        # 构建模型参数字典
        model_params = {}
        missing_fields = []
        
        for chinese_desc, model_field in field_mapping.items():
            value = form_data.get(chinese_desc, "").strip()
            if not value:
                # 检查是否为必填字段
                if model_field != "steel_density":  # steel_density有默认值，不是必填
                    missing_fields.append(chinese_desc)
                continue
                
            try:
                # 尝试转换为适当的类型
                if model_field == "well_name":
                    model_params[model_field] = value  # 字符串类型
                elif model_field in ["friction_coefficient"]:
                    model_params[model_field] = float(value)
                else:
                    model_params[model_field] = float(value)
            except ValueError:
                raise ValueError(f"字段 '{chinese_desc}' 的值 '{value}' 不是有效的数字")

        # 检查是否有必填字段缺失
        if missing_fields:
            missing_fields_str = "、".join(missing_fields)
            raise ValueError(f"以下必填参数不能为空：\n{missing_fields_str}")

        try:
            return DrillingParameters(**model_params)
        except Exception as e:
            # 解析Pydantic错误信息，提供更友好的提示
            error_msg = str(e)
            if "validation error" in error_msg:
                # 提取缺失的字段名
                if "Field required" in error_msg:
                    # 找出缺失的字段
                    for model_field, chinese_desc in {v: k for k, v in field_mapping.items()}.items():
                        if model_field in error_msg and model_field not in model_params:
                            raise ValueError(f"必填参数 '{chinese_desc}' 不能为空")
                raise ValueError(f"数据验证失败：{error_msg}")
            else:
                raise ValueError(f"创建钻井参数模型失败: {error_msg}")

    def set_form_data(self, data):
        """设置表单数据"""
        for param, field in self.input_fields.items():
            field.setText(data.get(param, ""))

    def clear_form(self):
        """清空表单"""
        for field in self.input_fields.values():
            field.clear()
        # 清空表格
        self.ui.tableWidget_zzsj.setRowCount(0)
        self.current_well_name = None
        self.data_modified = False
        self.update_status()

    def check_unsaved_changes(self):
        """检查是否有未保存的更改"""
        if self.data_modified:
            reply = QMessageBox.question(
                None, "未保存的更改",
                "当前数据尚未保存，是否保存?",
                QMessageBox.Save | QMessageBox.Discard | QMessageBox.Cancel
            )

            if reply == QMessageBox.Save:
                return self.save_well()
            elif reply == QMessageBox.Discard:
                return True
            elif reply == QMessageBox.Cancel:
                return False

        return True

    def new_well(self):
        """新建油井"""
        if not self.check_unsaved_changes():
            return

        self.clear_form()

    def manage_wells(self):
        """打开油井管理对话框"""
        if not self.check_unsaved_changes():
            return

        dialog = WellManagerDialog(self.main, self.well_data, self.current_well_name)
        if dialog.exec_() == QDialog.Accepted and dialog.action == 'open' and dialog.selected_well:
            well_name = dialog.selected_well
            well_data = self.well_data.get(well_name)

            if well_data:
                # 设置表单数据
                self.set_form_data(well_data)
                # 设置表格数据（如果有）
                table_data = well_data.get("table_data", [])
                self.set_table_data(table_data)
                
                self.current_well_name = well_name
                self.data_modified = False
                self.update_status()

        # 更新数据（可能在对话框中进行了重命名或删除操作）
        self.well_data = dialog.well_data
        self.save_data()

    def save_well(self):
        """保存油井数据"""
        form_data = self.get_form_data()
        # 获取表格数据
        table_data = self.get_table_data()
        
        # 检查油井名称
        well_name = form_data.get("井号", "").strip()
        if not well_name:
            QMessageBox.warning(None, "警告", "请输入井号!")
            return False
        try:
            self.get_drilling_parameters(form_data)
        except ValueError as e:
            QMessageBox.warning(None, "数据验证错误", str(e))
            return False

        # 将表格数据添加到表单数据中
        combined_data = form_data.copy()
        if table_data:
            combined_data["table_data"] = table_data

        # 检查是否已存在同名油井（排除当前正在编辑的油井）
        well_exists = well_name in self.well_data
        current_is_same = self.current_well_name == well_name

        # 如果当前没有正在编辑的油井，或者油井名称已更改
        if not self.current_well_name or not current_is_same:
            if well_exists:
                # 油井名称已存在，提示用户
                reply = QMessageBox.question(
                    None, "油井已存在",
                    f"油井 '{well_name}' 已存在，是否覆盖?",
                    QMessageBox.Yes | QMessageBox.No
                )
                if reply == QMessageBox.No:
                    return False
                # 如果选择覆盖，直接保存
                self.well_data[well_name] = combined_data
                if self.save_data():
                    self.current_well_name = well_name
                    self.data_modified = False
                    self.update_status()
                    QMessageBox.information(None, "成功", f"已修改油井数据: {well_name}")
                    return True
            else:
                # 新油井名称，直接保存
                self.well_data[well_name] = combined_data
                if self.save_data():
                    self.current_well_name = well_name
                    self.data_modified = False
                    self.update_status()
                    QMessageBox.information(None, "成功", f"已保存为新油井数据: {well_name}")
                    return True
        else:
            # 保存当前油井（名称未改变）
            if self.current_well_name and current_is_same:
                self.well_data[well_name] = combined_data
                if self.save_data():
                    self.data_modified = False
                    self.update_status()
                    QMessageBox.information(None, "成功", f"已保存油井数据: {well_name}")
                    return True

        return False